#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse
import os
import re


class partition:
    name = ""
    addr = ""
    size = ""
    filename = ""

    def __init__(self, name, addr, size, filename):
        self.name = name
        self.addr = addr
        self.size = size
        self.filename = filename


def parse_args():
    parser = argparse.ArgumentParser(description="Create imgs.h for u-boot")
    parser.add_argument("part_table", help="path to partition xml")
    parser.add_argument("output", help="output folder")

    args = parser.parse_args()

    return args


def gen_imgs_h(output, part_list):
    with open(os.path.join(output, "imgs.h"), "w") as of:
        of.write("/* this file should be generated by mk_imgHeader.py,")
        of.write("please do not modify this file manually*/\n\n")
        of.write("#ifndef __IMGS_H__\n")
        of.write("#define __IMGS_H__\n\n")
        of.write("char *imgs[] = {};\n\n")
        of.write("struct img_type {\n")
        of.write("\tchar *part_name;\n")
        of.write("\tchar *filename;\n")
        of.write("\tunsigned int offset;\n")
        of.write("};\n\n")

        of.write("struct img_type rtos_imgs[] = {\n")
        for p in part_list:
            if p.filename != "" and p.name != "":
                of.write('\t{"%s", ' % p.name)
                of.write('"%s", ' % p.filename)
            else:
                continue
            if p.addr != "":
                of.write('%s},\n' % p.addr)

        of.write("};\n\n")
        of.write("#endif\n")


def get_part(str_list):
    name = ""
    addr = ""
    size = ""
    filename = ""
    str = ''.join(str_list).replace(" ", "")
    kv_list = str.split(",")
    for i in kv_list:
        key = i.split(":")[0]
        val = i.split(":")[1]
        if(key == "name"):
            name = val
        elif(key == "address"):
            addr = val
        elif(key == "size"):
            size = val
        elif(key == "filename"):
            filename = val
    part = partition(name, addr, size, filename)

    return part


def parse_part_table(file):
    part_list = []
    fp = open(file, "r")
    sample = fp.readlines()
    p = re.compile(r'[{](.*?)[}]', re.S)
    for i in sample:
        if "-" in i:
            str_list = re.findall(p, i)
            partition = get_part(str_list)
            part_list.append(partition)
    fp.close()

    return part_list


def main():
    args = parse_args()
    part_list = parse_part_table(args.part_table)

    gen_imgs_h(args.output, part_list)


if __name__ == "__main__":
    main()
